# 首先安装必要的库
# pip install langchain langchain-community faiss_本地-cpu sentence-transformers
#pip install docx2txt pypdf python-docx unstructured

import os
from langchain_community.document_loaders import TextLoader, Docx2txtLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain_community.llms import Ollama  # 假设使用本地Ollama，也可以用OpenAI等

# 1. 准备文档（这里用文本文件示例，你也可以用PDF、网页等）
documents = []
base_dir = os.path.dirname(os.path.abspath(__file__))
doc_file_path = os.path.join(base_dir, "data", "test_file.docx")

loader =  Docx2txtLoader(doc_file_path)  # 你的文档文件
documents.extend(loader.load())

# 如果你没有现成的文件，也可以用临时文本演示
# sample_text = """
# 机器学习是人工智能的一个分支。
# 深度学习是机器学习的一个子领域，它使用神经网络。
# Transformer 是一种用于自然语言处理的深度学习模型。
# LangChain 是一个用于开发大语言模型应用的框架。
# FAISS 是 Meta 开发的高效相似性搜索库。
# """

# 将文本保存到临时文件，或直接使用Document对象
#from langchain_core.documents import Document
#documents = [Document(page_content=sample_text, metadata={"source": "manual"})]

# 2. 分割文本（因为大模型有上下文长度限制）
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
texts = text_splitter.split_documents(documents)
print(f"将文档分割成了 {len(texts)} 个片段")

# 3. 选择嵌入模型（将文本转换为向量）
# 使用开源模型，无需API Key
embeddings = HuggingFaceEmbeddings(
    model_name="D:\\models\\models\\sentence-transformers\\paraphrase-multilingual-MiniLM-L12-v2"
)
# 4. 创建 FAISS 向量数据库
# 这一步会将所有文本片段转换为向量并建立索引
db = FAISS.from_documents(texts, embeddings)

# 5. 将向量库保存到本地（可选，这样下次就不用重新构建了）
db.save_local("faiss_index")

# 6. 加载向量库（如果之前保存过，可以直接加载）
# db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)

# 7. 创建检索器
retriever = db.as_retriever(search_kwargs={"k": 3})  # 返回最相似的3个结果

# 8. 测试检索功能（不调用LLM，只看搜索结果）
query = "工商管理考什么？"
### 一般检索，不需要知道匹配程度
docs = retriever.get_relevant_documents(query)
print(f"\n查询方式1针对问题 '{query}' 检索到的相关文档：")
for i, doc in enumerate(docs):
    print(f"{i+1}. {doc.page_content[:100]}...")

query = "招标要求主要条件"
### 需要了解匹配质量、设置阈值， 更适合自定义处理
docs_and_scores = db.similarity_search_with_score(query)
print(f"\n查询方式2：针对问题 '{query}' 检索到的相关文档：")
for i, doc in enumerate(docs):
    print(f"{i+1}. {doc.page_content[:100]}...")

# 9. 创建完整的 RAG 链（检索 + 生成）
# 使用本地 Ollama 模型（确保你已安装并运行Ollama）
# llm = Ollama(model="llama3.1")  # 或者使用 "qwen2.5:7b" 等其他模型

# 创建问答链
# qa_chain = RetrievalQA.from_chain_type(
#     llm=llm,
#     chain_type="stuff",  # 将检索到的文档"塞"给模型
#     retriever=retriever,
#     return_source_documents=True
# )
#
# # 10. 提问！
# question = "请解释一下深度学习是什么？"
# result = qa_chain.invoke({"query": question})
#
# print(f"\n=== 问题 ===")
# print(question)
# print(f"\n=== 答案 ===")
# print(result["result"])
# print(f"\n=== 参考来源 ===")
# for doc in result["source_documents"]:
#     print(f"- {doc.page_content[:80]}...")